from pyecharts import options as opts
from pyecharts.charts import Graph
import os

class KeywordBubbleChart:
    """核心词共现关系气泡图生成器"""
    
    def __init__(self):
        """
        初始化颜色配置
        """
        self.colors = [
            "#E0FFFF", "#FFFFF0", "#D8BFD8", "#FFEFD5",
            "#7FFFD4", "#FFC0CB", "#ADD8E6", "#FAFAD2"
        ]
    
    def generate(self, keywords, word_pairs, save_path, max_size=70, min_size=25):
        """生成气泡图
        
        Args:
            keywords: Dict[str, int] - 核心词及其出现次数
            word_pairs: List[Tuple] - 词对关系列表 [(word1, word2, co_count),...]
            save_path: str - 保存路径
            max_size: int - 最大气泡大小
            min_size: int - 最小气泡大小
        """
        # 获取最大出现次数
        max_count = max(keywords.values())
        
        # 创建节点数据
        nodes = [
            {
                "name": word,
                "symbolSize": max(min_size, (count / max_count) * max_size),  # 根据比例调整气泡大小
                "value": count,
                "category": i % len(self.colors),  # 确保 category 值在 colors 范围内
                "label": {
                    "show": True,
                    "position": "inside",
                    "fontSize": 12,  # 调整字体大小
                    "color": "#333",  # 字体颜色
                    "fontWeight": "bold"  # 加粗字体
                },
                "itemStyle": {
                    "opacity": 0.95  # 设置气泡透明度
                }
            }
            for i, (word, count) in enumerate(keywords.items())
        ]
        
        # 创建 categories 配置
        categories = [
            {"name": str(i), "itemStyle": {"color": self.colors[i]}}  # 将颜色与 category 关联
            for i in range(len(self.colors))
        ]
        
        # 创建边数据，并根据共现次数动态设置线段粗细
        max_co_count = max(co_count for _, _, co_count in word_pairs)
        links = [
            {
                "source": word1,
                "target": word2,
                "value": co_count,
                "lineStyle": {
                    "width": (co_count / max_co_count) * 5,  # 减小连线宽度
                    "color": "rgba(100, 100, 100, 0.5)"  # 使用半透明的灰色
                }
            }
            for word1, word2, co_count in word_pairs
        ]
        
        # 创建图实例
        chart = Graph(
            init_opts=opts.InitOpts(
                width="600px",  # 增加画布宽度
                height="400px",  # 增加画布高度
                bg_color="#ffffff"
            )
        )
        
        # 添加数据
        chart.add(
            series_name="",
            nodes=nodes,
            links=links,
            categories=categories,  # 使用配置好的 categories
            layout="force",
            is_roam=True,
            is_focusnode=True,
            repulsion=200,  # 增加节点之间的排斥力
            edge_length=170,  # 增加边的长度
            label_opts=opts.LabelOpts(
                position="inside",
                font_size=12  # 调整标签字体大小
            ),
            tooltip_opts=opts.TooltipOpts(
                trigger="item",
                formatter="{b}: {c}"
            ),
        )
        
        # 设置全局配置
        chart.set_global_opts(
            title_opts=opts.TitleOpts(
                is_show=False,
                title="核心词共现关系图",
                subtitle="气泡大小表示词频，连线粗细表示共现次数",
                pos_left="center",
                pos_top="10px",  # 将标题和描述向上移动
                title_textstyle_opts=opts.TextStyleOpts(font_size=14),  # 调整标题字体大小
                subtitle_textstyle_opts=opts.TextStyleOpts(font_size=12)  # 调整副标题字体大小
            ),
            legend_opts=opts.LegendOpts(is_show=False),
        )
        
        # 生成图表
        os.makedirs(os.path.dirname(save_path), exist_ok=True)
        chart.render(save_path)
        return save_path